import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from PIL import Image

# # 确保 TensorFlow 使用 GPU
# gpus = tf.config.experimental.list_physical_devices('GPU')
# if gpus:
#     try:
#         # 设置 GPU 显存动态增长，避免一次性占满
#         for gpu in gpus:
#             tf.config.experimental.set_memory_growth(gpu, True)
#         tf.config.experimental.set_virtual_device_configuration(
#             gpus[0],
#             [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)])  # 设置显存限制为4GB，根据你的需要调整
#     except RuntimeError as e:
#         print(e)
# # 从 TensorFlow Hub 加载 ESRGAN 模型
# model = hub.load("https://tfhub.dev/captain-pool/esrgan-tf2/1")
#
# # 加载低分辨率图像并调整大小
# image = Image.open('./9322cc7777230e3032ead7add263974.png').convert("RGB")
# image = np.array(image)
#
# # 将图像尺寸调整到模型所需的输入尺寸
# image = tf.image.resize(image, [tf.cast(tf.shape(image)[0] / 4, tf.int32),
#                                 tf.cast(tf.shape(image)[1] / 4, tf.int32)])
#
# # 将图像转换为 Tensor 并添加批次维度
# image = tf.convert_to_tensor(image, dtype=tf.float32)
# image = tf.expand_dims(image, axis=0)
#
# # 使用模型生成超分辨率图像
# upscaled_image = model(image)
#
# # 移除批次维度并转换回 numpy 数组
# upscaled_image = tf.squeeze(upscaled_image)
# upscaled_image = tf.clip_by_value(upscaled_image, 0, 255)
# upscaled_image = tf.cast(upscaled_image, tf.uint8).numpy()
#
# # 保存或显示结果图像
# upscaled_image_pil = Image.fromarray(upscaled_image)
# upscaled_image_pil.save("upscaled_image.png")
# # upscaled_image_pil.show()

import tensorflow as tf

# 打印 TensorFlow 可用的物理设备
print("Available devices:", tf.config.experimental.list_physical_devices())

# 检查 TensorFlow 是否使用 GPU
if tf.config.experimental.list_physical_devices('GPU'):
    print("TensorFlow is using GPU")
else:
    print("TensorFlow is using CPU")
